package com.piece.core.oauth.handler;

import com.piece.core.framework.constant.ExceptionConstants;
import com.piece.core.log.config.TenantContextHolder;
import com.piece.core.framework.constant.ExceptionAuthConstants;
import com.piece.core.framework.constant.SecurityConstants;
import com.piece.core.framework.support.response.AjaxResponse;
import com.piece.core.framework.util.basic.I18nUtil;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.httpclient.HttpStatus;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.springframework.http.ResponseEntity;
import org.springframework.security.authentication.InsufficientAuthenticationException;
import org.springframework.security.core.Authentication;
import org.springframework.security.oauth2.common.OAuth2AccessToken;
import org.springframework.security.oauth2.common.util.OAuth2Utils;
import org.springframework.security.oauth2.provider.OAuth2Authentication;
import java.security.Principal;
import java.util.Map;

/**
 * oauth2登录切面
 */
@Slf4j
@Aspect
public class TokenEndpointHandler {

    @Around("execution(* org.springframework.security.oauth2.provider.endpoint.TokenEndpoint.postAccessToken(..))")
    public Object handleControllerMethod(ProceedingJoinPoint joinPoint) throws Throwable {
        try {
            Object[] args = joinPoint.getArgs();
            Principal principal = (Principal) args[0];
            if (!(principal instanceof Authentication)) {
                throw new InsufficientAuthenticationException(
                        "There is no client authentication. Try adding an appropriate authentication filter.");
            }

            String clientId = getClientId(principal);
            TenantContextHolder.setTenant(clientId);

            Map<String, String> parameters = (Map<String, String>) args[1];
            String grantType = parameters.get(OAuth2Utils.GRANT_TYPE);

            Object proceed = joinPoint.proceed(args);
            if (SecurityConstants.AUTHORIZATION_CODE.equals(grantType)) {
                return proceed;
            }
            if (null != proceed) {
                ResponseEntity<OAuth2AccessToken> responseEntity = (ResponseEntity<OAuth2AccessToken>) proceed;
                OAuth2AccessToken body = responseEntity.getBody();
                if (responseEntity.getStatusCode().is2xxSuccessful()) {
                    return ResponseEntity
                            .status(HttpStatus.SC_OK)
                            .body(AjaxResponse.success(I18nUtil.message(ExceptionConstants.MESSAGE_SUCCESS), body));
                }
            }
            return ResponseEntity
                    .status(HttpStatus.SC_OK)
                    .body(AjaxResponse.error(I18nUtil.message(ExceptionAuthConstants.AUTHENTICATION)));
        } finally {
            TenantContextHolder.clear();
        }
    }

    private String getClientId(Principal principal) {
        Authentication client = (Authentication) principal;
        if (!client.isAuthenticated()) {
            throw new InsufficientAuthenticationException("The client is not authenticated.");
        }
        String clientId = client.getName();
        if (client instanceof OAuth2Authentication) {
            clientId = ((OAuth2Authentication) client).getOAuth2Request().getClientId();
        }
        return clientId;
    }
}
